
from collections import deque
from threading import Thread, Lock
from multiprocessing import Process
import time, random, os
from DataLoader import DataLoader
from MQTTTools import sendMQTT_transport_car, sendMQTT_car_count
from .MQTTData.VehicleDetecteMsg import *
import torch
from ReidTools import getFeatureMap, getFeatureMaps, getFeatureMapsToFMList, getFeatureMapsFromDumpsToFMList, getFeatureMapFromArray
import pickle


class DataProcessor:

    def __init__(self, parameters, commonMemory,
                 taskScheduler, dataLoader:DataLoader):
        '''
        :param parameters: 工程参数
        :param CommonMemory: 共享内存对象
        :param TaskScheduler: 任务调度器对象
        :param DataLoader: 数据加载器对象
        '''

        # 工程参数
        self.parameters = parameters
        # 线程锁
        self.lock = Lock()
        # 共享内存对象
        self.cm = commonMemory
        # 任务调度器对象
        self.taskScheduler = taskScheduler
        # 数据加载器对象
        self.dataLoader = dataLoader

        # 储存CameraLists, CameraLists是一个字典，每个键都对应一个Camera列表
        self.cameraLists = parameters["CameraLists"]
        # 每个路段的摄像头储存队列
        self.dataQueuesDict = {}
        # 每个路段最前一个摄像头的特征池，用于被查询
        self.featureMapsDict = {}
        # 每个路段的每个储存队列最多储存多少个信息
        self.queuesMaxLeng = {}
        # 获取到消息的记录
        self.cameraGetDataRecord = {k:torch.tensor([0] * len(self.cameraLists[k]), dtype=torch.int8)
                                    for k in self.cameraLists}
        # 每个路段有多少个摄像头
        self.oneRoadCameraNum = {k:len(self.cameraLists[k]) for k in self.cameraLists}

        for road in self.cameraLists:
            self.dataQueuesDict[road] = {}
            for camera in self.cameraLists[road]:
                # 为当前路段每个摄像头都设置储存队列
                self.dataQueuesDict[road][camera] = deque()
            # 按公式设置当前路段每个储存队列最多储存多少个信息
            self.queuesMaxLeng[road] = (len(self.cameraLists[road]) + 1) * 6
            # self.queuesMaxLeng[road] = 5
            # 为每个reid特征池建立队列
            self.featureMapsDict[road] = deque()

        from ReidTools.FeatureMapTools import getFeatureMap
        self.getFeatureMap = getFeatureMap
        self.model = torch.load("ReidTools/model_cpu.pt").to(torch.device("cuda"))
        self.model.eval()
        # print(self.model)
        print("load model over！")

    def start(self):
        # 开启读取和发送线程
        # readThread = Thread(target=self._readDataThread)
        # readThread.start()
        self._readDataThread()

        # sendThread = []
        # for road in self.cameraLists:
        #     t = Thread(target=self._oneRoadSendDataThread, args=(road, ))
        #     sendThread.append(t)
        #     t.start()

        # readThread = Process(target=self._readDataThread)
        # readThread.start()
        #
        # sendThread = []
        # for road in self.cameraLists:
        #     t = Process(target=self._oneRoadSendDataThread, args=(road, ))
        #     sendThread.append(t)
        #     t.start()



    def _readDataThread(self):
        while True:
            # 如果读到重整信号，所有储存队列清零
            if self.cm.isReform():
                print("读取到重整信号！3秒后开始重整！")
                time.sleep(3)
                with self.lock:
                    for k in self.cameraLists:
                        self.dataQueuesDict[k] = {}
                        for camera in self.cameraLists[k]:
                            # 为当前路段每个摄像头都设置储存队列
                            # if self.dataQueuesDict[k][camera]:
                            #     self.dataQueuesDict[k][camera].pop()
                            self.dataQueuesDict[k][camera] = deque()
                        # 按公式设置当前路段每个储存队列最多储存多少个信息
                        self.dataQueuesDict[k]["queuesMaxLeng"] = (len(self.cameraLists[k]) + 1) * 6
                        if self.featureMapsDict[k]:
                            self.featureMapsDict[k].pop()

                print("重整完毕！")
                self.cm.setWaitingFlag()

            if self.cm.canRead():
                self.cm.setReadingFlag()
                # data = cm._mem[1:(cm._mem[3] * 8 + 4)]
                data = [self.cm.getData(i) for i in range(self.cm._mem[3] * 8 + 5)]
                self.cm.setWaitingFlag()

                # 获取摄像头id
                camera = data[1]
                road = camera.split("+")[0]

                # 计算FeatureMap
                # FeatureMapSaveList = []
                # for i in range(5, len(data), 8):
                #
                #     # t = Thread(target=getFeatureMap, args=(self.taskScheduler, self.parameters["torchDevice"],
                #     #                                        data[i], os.path.dirname(data[i])))
                #     # t.start()
                #     # t = Thread(target=getFeatureMap, args=(self.model, self.parameters["torchDevice"],
                #     #                                        data[i], os.path.dirname(data[i])))
                #     # t.start()
                #     getFeatureMap(self.model, self.parameters["torchDevice"],
                #                   data[i], os.path.dirname(data[i]))
                # if len(data) > 5:
                #     getFeatureMaps(self.model, self.parameters["torchDevice"], os.path.dirname(data[5]))

                # fm = getFeatureMapsToFMList(self.model, self.parameters["torchDevice"], os.path.dirname(data[5])) \
                #     if len(data) > 5 else []

                fm = getFeatureMapsFromDumpsToFMList(
                    self.model, self.parameters["torchDevice"], [data[i] for i in range(5, len(data), 8)]
                ) if len(data) > 5 else []

                with self.lock:
                    # 添加保存记录
                    self.cameraGetDataRecord[road][self.cameraLists[road].index(camera)] = 1
                    # 将整个data都保存在队列中
                    self.dataQueuesDict[road][camera].append(data)
                    if self.dataQueuesDict[road][camera].__len__() > self.queuesMaxLeng[road]:
                        self.dataQueuesDict[road][camera].popleft()
                    print("camera ", camera, ": ", self.dataQueuesDict[road][camera].__len__(), "个数据")

                    # # 获取特征列表
                    # fmList = [ [int(data[i].split("/")[-1][3:-4]),
                    #             data[i][:-3] + "pt"]
                    #            for i in range(5, len(data), 8)]
                    # # 预加载特征文件
                    # self.dataLoader.reload(fmList)

                    # 如果是当前路段最前面的摄像头，则需要储存特征路径到内存里的队列
                    if camera == self.cameraLists[road][0]:
                        # self.featureMapsDict[road].append(fmList)
                        self.featureMapsDict[road].append(fm)

                        if self.featureMapsDict[road].__len__() > self.queuesMaxLeng[road]:
                            # 弹出不需要的数据的同时，清理ReIDTools里辅助加载特征文件字典里不需要的Tensor
                            self.featureMapsDict[road].popleft()
                            # l = self.featureMapsDict[road].popleft()
                            # self.dataLoader.cleanLoadFeatureMapDict(l)

                    # print("放入数据结束！")
                self._oneRoadOnceSendData(road)


    def _oneRoadOnceSendData(self, road:str):

        if int(self.cameraGetDataRecord[road].sum(0)) != self.oneRoadCameraNum[road]:
            return
        else:
            self.cameraGetDataRecord[road] -= 1

        cameraList = self.cameraLists[road]
        dataQueues = self.dataQueuesDict[road]
        dataQueuesMaxLeng = self.queuesMaxLeng[road]
        featureMapPaths = self.featureMapsDict[road]

        if True:
            # 如果队列未满则跳过
            if dataQueues[cameraList[-1]].__len__() < dataQueuesMaxLeng:
                # print(dataQueues[cameraList[-1]].__len__() )
                return

            # 获取开始时间
            startTime = time.time()

            # 要准备发送的各摄像头在队列中的data
            tempToBeSendData = []
            # 拷贝准备比对的特征集合
            # tempFeatureMap = []
            # tempFeatureMapIdx = []
            with self.lock:
                for camera in cameraList:
                    # data = dataQueues[camera].pop()
                    # dataQueues[camera].append(data)
                    data = dataQueues[camera][-1]
                    tempToBeSendData.append(data)
                tempFeatureMapWithIdx = [i for l in featureMapPaths for i in l]

            # 加载被查询池特征，此处为路段第一个摄像头的所有队列的特征
            # queriedPoolFeatureMaps = torch.stack([self.dataLoader.loadFeatureMap(i[1])
            #                               for i in tempFeatureMapWithIdx], dim=0)
            queriedPoolFeatureMaps = torch.stack([i[1] for i in tempFeatureMapWithIdx], dim=0)

            # 对每个摄像头发送数据
            for cameraListIdx in range(len(cameraList)):

                # 构建MQTT数据类
                camera = cameraList[cameraListIdx]

                code = 0
                timestamp = int(time.time())
                msg = u"成功"

                roadpart = camera.split("+")[0][1:]
                roadpartx, roadparty = 0, 0
                if road == "K35":
                    direction = 2
                elif road == "K28":
                    direction = 1
                else:
                    direction = 0

                # 获取当前摄像头的数据并开始整理
                data = tempToBeSendData[cameraListIdx]

                vdr = VehicleDetectResult(roadpart, roadpartx, roadparty, direction, camera)
                vdcc = VehicleDetectCarCount(roadpart, roadpartx, roadparty, direction, ((data[3] // 2) + (data[3] % 2)))

                for i in range(5, 5 + (data[3] * 8), 8):
                    data[i] = pickle.loads(data[i])
                    # 获取数据
                    num, carid, cartype, carcolor, roadlane, distance, curx, cury = data[i: i + 8]

                    # 由于坐标原点在左上角，需要转化一下
                    cury = self.parameters["PrimerImgHigh"] - cury
                    # 将像素转化为米
                    curx *= self.parameters["pixelScale"][camera]["x"]
                    cury *= self.parameters["pixelScale"][camera]["y"]
                    cury += self.parameters["mileageBios"][camera]["x"]
                    # distance = cury * 2  # 转化后扩展
                    distance = float((curx ** 2 + cury ** 2) ** 0.5)

                    # 车牌号暂时无法识别，留空
                    carid = ""

                    # 颜色string转int，没有的颜色赋值为红色
                    # _colorType = {"red": 1, "orange": 2, "yellow": 3, "green": 4, "cyan": 5, "blue": 6, "purple": 7}
                    _colorType = {"red": 1, "blue": 2, "green": 3, "yellow": 4, "black": 5, "white": 6, "gray": 7, "purple":8}
                    # carcolor = (lambda colorStr: _colorType[colorStr] if colorStr in _colorType else 1)(carcolor)
                    if carcolor in _colorType.keys():
                        carcolor = _colorType[carcolor]
                    else:
                        carcolor = 1

                    # 此时num里填的是车辆图片文件路径字符串，文件名为id_xx.jpg，xx为MOT生成的num
                    # 如果是该路段首个摄像头，则直接将MOT生成的ID作为num，否则启动reid为它们重写num值
                    if cameraListIdx == 0:
                        # num = data[i].split("/")[-1][3:-4]
                        num = data[i][0]
                    else:
                        # reid
                        # # 加载query图片特征 version 1
                        # queryFeatureMapPath = num[:-3] + "pt"
                        # while not os.path.exists(queryFeatureMapPath):
                        #     time.sleep(0.01)
                        # queryFeatureMap = self.dataLoader.torchLoad(queryFeatureMapPath)
                        # # 加载query图片特征 version 2
                        # queryFeatureMap = getFeatureMap(self.model, "cuda", imgPath=num, savePath=None, hasRet=True)
                        # 加载query图片特征 version 3
                        queryFeatureMap = getFeatureMapFromArray(self.model, "cuda", imgArray=data[i][1])
                        # query和tempFeatureMap里的所有特征计算距离
                        minDistanceIdx = (queriedPoolFeatureMaps - queryFeatureMap).pow(2).sum(1).argmin()
                        num = tempFeatureMapWithIdx[int(minDistanceIdx)][0]
                    # 添加车辆数据
                    vdr.addCarInfo(num=num, carid=carid, cartype=cartype, carcolor=carcolor, distance=distance,
                                   curx=curx, cury=cury, roadlane=roadlane)

                vdm_result = str(VehicleDetectMsg(code, timestamp, msg, vdr).dict)
                vdm_carCount = str(VehicleDetectMsg(code, timestamp, msg, vdcc).dict)
                # print("发送数据：", vdm_result)

                # t = Thread(target=sendMQTTThread, args=(vdm, times, sleepNum))
                # t.start()
                sendMQTT_transport_car(vdm_result)
                sendMQTT_car_count(vdm_carCount)

            # sleepNum = 20 / 24 - (time.time() - startTime)
            # if sleepNum > 0:
            #     time.sleep(sleepNum)
            #     print("sleep: ", sleepNum)

    # def _oneRoadSendDataThread(self, road:str):
    #
    #     cameraList = self.cameraLists[road]
    #     dataQueues = self.dataQueuesDict[road]
    #     dataQueuesMaxLeng = self.queuesMaxLeng[road]
    #     featureMapPaths = self.featureMapsDict[road]
    #
    #     if True:
    #         # 如果队列未满则跳过
    #         if dataQueues[cameraList[-1]].__len__() < dataQueuesMaxLeng:
    #             # print(dataQueues[cameraList[-1]].__len__() )
    #             return
    #         # 要准备发送的各摄像头在队列中的data
    #         tempToBeSendData = []
    #         # 拷贝准备比对的特征集合
    #         # tempFeatureMap = []
    #         # tempFeatureMapIdx = []
    #         with self.lock:
    #             for camera in cameraList:
    #                 # data = dataQueues[camera].pop()
    #                 # dataQueues[camera].append(data)
    #                 data = dataQueues[camera][-1]
    #                 tempToBeSendData.append(data)
    #             tempFeatureMapWithIdx = [i for l in featureMapPaths for i in l]
    #
    #         # 加载特征
    #         # for i in tempFeatureMapWithIdx:
    #         #     ten = self.dataLoader.loadFeatureMap(i[1])
    #         #     print(i, ten.shape)
    #         tempFeatureMap = torch.stack([self.dataLoader.loadFeatureMap(i[1])
    #                                       for i in tempFeatureMapWithIdx], dim=0)
    #         # 对每个摄像头发送数据
    #         for cameraListIdx in range(len(cameraList)):
    #             print(cameraListIdx)
    #
    #             # 构建MQTT数据类
    #             camera = cameraList[cameraListIdx]
    #
    #             code = 0
    #             timestamp = int(time.time())
    #             msg = u"成功"
    #
    #             roadpart = camera.split("+")[0][1:]
    #             roadpartx, roadparty = 0, 0
    #             if road == "K35":
    #                 direction = 2
    #             elif road == "K28":
    #                 direction = 1
    #             else:
    #                 direction = 0
    #
    #             vdr = VehicleDetectResult(roadpart, roadpartx, roadparty, direction, camera)
    #
    #             # 获取当前摄像头的数据并开始整理
    #             data = tempToBeSendData[cameraListIdx]
    #             for i in range(5, 5 + (data[3] * 8), 8):
    #
    #                 # 获取数据
    #                 num, carid, cartype, carcolor, roadlane, distance, curx, cury = data[i: i + 8]
    #
    #                 # 由于坐标原点在左上角，需要转化一下
    #                 cury = self.parameters["PrimerImgHigh"] - cury
    #                 # distance = cury * 2  # 转化后扩展
    #                 distance = float((curx ** 2 + cury ** 2) ** 0.5)
    #
    #                 # 车牌号暂时无法识别，留空
    #                 carid = ""
    #
    #                 # 颜色string转int，没有的颜色赋值为红色
    #                 _colorType = {"red": 1, "blue": 2, "green": 3, "yellow": 4, "black": 5, "white": 6, "gray": 7, "purple":8}
    #                 # carcolor = (lambda colorStr: _colorType[colorStr] if colorStr in _colorType else 1)(carcolor)
    #                 if carcolor in _colorType.keys():
    #                     carcolor = _colorType[carcolor]
    #                 else:
    #                     carcolor = 1
    #
    #                 # 此时num里填的是车辆图片文件路径字符串，文件名为id_xx.jpg，xx为MOT生成的num
    #                 # 如果是该路段首个摄像头，则直接将MOT生成的ID作为num，否则启动reid为它们重写num值
    #                 if cameraListIdx == 0:
    #                     num = data[i].split("/")[-1][3:-4]
    #                 else:
    #                     # reid
    #                     # 加载query图片特征
    #                     queryFeatureMapPath = num[:-3] + "pt"
    #                     while not os.path.exists(queryFeatureMapPath):
    #                         time.sleep(0.01)
    #                     queryFeatureMap = self.dataLoader.torchLoad(queryFeatureMapPath)
    #                     # query和tempFeatureMap里的所有特征计算距离
    #                     # if True:
    #                     #     def featureMapDistance(queryTensor, cmpTensor):
    #                     #         distance = (cmpTensor - queryTensor).pow(2).sum(0)
    #                     #         distance = float(distance)
    #                     #         return distance
    #                     #     featureMapCompareRes = [(fmData[0], featureMapDistance(queryFeatureMap, fmData[1]))
    #                     #                             for fmData in tempFeatureMap]
    #                     #     # 排序
    #                     #     featureMapCompareRes = sorted(featureMapCompareRes, key=lambda x: x[1])
    #                     #     # print(featureMapCompareRes)
    #                     #     num = featureMapCompareRes[0][0]
    #
    #                     minDistanceIdx = (tempFeatureMap - queryFeatureMap).pow(2).sum(1).argmin()
    #                     num = tempFeatureMapWithIdx[int(minDistanceIdx)][0]
    #
    #                 # 添加车辆数据
    #                 vdr.addCarInfo(num=num, carid=carid, cartype=cartype, carcolor=carcolor, distance=distance,
    #                                curx=curx, cury=cury, roadlane=roadlane)
    #
    #             vdm = VehicleDetectMsg(code, timestamp, msg, vdr)
    #             vdm = str(vdm.dict)
    #             print("发送数据：", vdm)
    #
    #             # t = Thread(target=sendMQTTThread, args=(vdm, times, sleepNum))
    #             # t.start()
    #             sendMQTT_transport_car(vdm)
    #
    #         # time.sleep(2)

